#include "manager/manager.h"

namespace infer {

TxnnManager::TxnnManager(const std::string& engine_path, const InferConfig& cfg)
  : engine_path_(engine_path), cfg_(cfg) {}

TxnnManager::TxnnManager() {}

void TxnnManager::init(const void* blob, std::size_t size) {
  if (!blob) {
    throw std::runtime_error("Failed to get a valid blob from param");
  }
  printf("init: blob size = %zu\n", size);
  // TODO 待添加 初始化 logger_
  // TODO 待添加 初始化 pluging 库兼容机制, 第二版重构

  // 创建 TxnnCpp Runtime
  // TODO 待添加 createTxnnCppRuntime(logger_); 直接调用 TxDA 初始化 runtime
  runtime_ = std::unique_ptr<ITMRuntime>();
  if (!runtime_) {
    throw std::runtime_error("Failed to create TxnnCpp Runtime");
  }

  // 反序列化引擎
  // TODO 待添加 runtime_->deserializeEngine(blob, size); 直接调用 TxDA 加载相关的编译产物
  engine_ = std::shared_ptr<ITMEngine>();
  if (!engine_) {
    throw std::runtime_error("Failed to deserialize engine");
  }

  // 创建执行上下文
  // TODO 待添加 engine_->createExecutionContext(); 直接调用 TxDA 初始化Context
  context_ = std::unique_ptr<ITMExecutionContext>();
  if (!context_) {
    throw std::runtime_error("Failed to create execution context");
  }
}

TxnnManager::~TxnnManager() {
  context_.reset();
  engine_.reset();
  runtime_.reset();
}

bool TxnnManager::setTensorAddress(const char* tensor_name, void* data) {
  if (!data) {
    throw std::runtime_error("Failed to get a valid data pointer from input_param");
  }
  // TODO 直接设置 TxDA 所需的 Tensor 输入数据
  printf("setTensorAddress: %s\n", tensor_name);
  return true;
}

bool TxnnManager::setInputShape(const char* tensor_name, const std::vector<dim_t>& dims) {
  if (dims.empty()) {
    throw std::runtime_error("Failed to get a valid dims from input_param");
  }
  printf("setInputShape: %s\n", tensor_name);
  // TODO 直接设置 TxDA 所需的 Tensor 输入形状

  return true;
}

bool TxnnManager::enqueue() {
  // TODO 直接调用 TxDA 执行推理
  return true;
}

std::vector<dim_t> TxnnManager::getTensorShape(const char* tensor_name) const noexcept {
  // TODO 直接获取 TxDA 所需的 Tensor 输出形状
  printf("getTensorShape: %s\n", tensor_name);
  return {};
}

std::int32_t TxnnManager::getTensorIOMode(const char* tensor_name) const noexcept {
  // TODO 直接获取 TxDA 所需的 Tensor 输入输出模式
  printf("getTensorIOMode: %s\n", tensor_name);
  return 1;
}

std::int32_t TxnnManager::getNbIOTensors() const noexcept {
  // TODO 直接获取 TxDA 所需的 Tensor 数量
  return 1;
}

const char* TxnnManager::getIOTensorName(std::int32_t index) const noexcept {
  printf("getIOTensorName: IOTensor index - %d\n", index);
  // TODO 直接获取 TxDA 所需的 Tensor 名称
  return "test";
}

} // namespace infer